/*
 *  jOTR - Java Off-The-Record Messaging Library
 *  Copyright (C) 2007-2008 Markus Karnik
 *
 *  This program is free software; you can redistribute it and/or modify it 
 *  under the terms of the GNU General Public License as published by the
 *  Free Software Foundation; either version 3 of the License, 
 *  or (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful, but
 *  WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
 *  or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License 
 *  for more details.
 *
 *  You should have received a copy of the GNU General Public License along 
 *  with this program; if not, see <http://www.gnu.org/licenses/>.
 *  
 */
package de.karnik.libjotr.helper;

import java.math.BigInteger;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;

import javax.crypto.Cipher;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import org.bouncycastle.crypto.params.DSAParameters;
import org.bouncycastle.crypto.params.DSAPrivateKeyParameters;
import org.bouncycastle.crypto.params.DSAPublicKeyParameters;
import org.bouncycastle.crypto.signers.DSASigner;

import de.karnik.helper.ArrayFunctions;
import de.karnik.helper.HashFunctions;
import de.karnik.helper.LogFunctions;
import de.karnik.helper.StringFunctions;
import de.karnik.libjotr.JotrConstants;
import de.karnik.libjotr.JotrContext;
import de.karnik.libjotr.JotrDsaKeyStorage;
import de.karnik.libjotr.JotrMpi;

/**
 * Common functions of the jOTR library. 
 * 
 * @author Markus Karnik - markus.karnik@my-design.org
 * @version 1.0
 * @since 1.0
 */
public abstract class JotrFunctions  implements JotrConstants {
	
	/**
	 * Constructor. Cannot be called.
	 */
	private JotrFunctions(){}
	
	/**
	 * Returns an encrypting AES-CTR-Cipher with length of keyBytes and initial counter value 0.
	 * 
	 * @param keyBytes The keyBytes for the secret key.
	 * @return An encrypting AES-CTR-Cipher with length of keyBytes and initial counter value 0.
	 * 
	 * @throws NoSuchPaddingException 
	 * @throws NoSuchAlgorithmException 
	 * @throws InvalidKeyException 
	 * @throws InvalidAlgorithmParameterException 
	 * @throws NoSuchProviderException 
	 */
	public static Cipher generateEncryptingAesCipher( byte[] keyBytes ) throws  NoSuchAlgorithmException, 
																				NoSuchPaddingException, 
																				InvalidKeyException, 
																				InvalidAlgorithmParameterException, 
																				NoSuchProviderException {
		return generateEncryptingAesCipher( keyBytes, CTR_ZERO_VAL );
	}
	
	/**
	 * Returns an encrypting AES-CTR-Cipher with length of keyBytes and initial counter value ctr.
	 * 
	 * @param keyBytes The keyBytes for the secret key.
	 * @param ctr Counter value as 8 bytes data. ( big endian )
	 * @return An encrypting AES-CTR-Cipher with length of keyBytes and initial counter value ctr.
	 * 
	 * @throws NoSuchAlgorithmException
	 * @throws NoSuchPaddingException
	 * @throws InvalidKeyException
	 * @throws InvalidAlgorithmParameterException
	 * @throws NoSuchProviderException
	 */
	public static Cipher generateEncryptingAesCipher( byte[] keyBytes, byte[] ctr ) throws  NoSuchAlgorithmException, 
																							NoSuchPaddingException, 
																							InvalidKeyException, 
																							InvalidAlgorithmParameterException,
																							NoSuchProviderException {
		// complete ctr key
		ctr = ArrayFunctions.merge( ctr, CTR_ZERO_VAL );
		
		IvParameterSpec spec = new IvParameterSpec( ctr );
		SecretKeySpec key = new SecretKeySpec( keyBytes, "AES" );
	    Cipher cipher = Cipher.getInstance( "AES/CTR/NoPadding" );

	    // encryption pass
	    cipher.init( Cipher.ENCRYPT_MODE, key, spec );
	    
		return cipher;
	}
	
	/**
	 * Returns a decrypting AES-CTR-Cipher with length of keyBytes and initial counter value 0.
	 * 
	 * @param keyBytes The keyBytes for the secret key.
	 * @return A decrypting AES-CTR-Cipher with length of keyBytes and initial counter value 0.
	 * 
	 * @throws NoSuchPaddingException 
	 * @throws NoSuchAlgorithmException 
	 * @throws InvalidKeyException 
	 * @throws InvalidAlgorithmParameterException 
	 * @throws NoSuchProviderException 
	 */
	public static Cipher generateDecryptingAesCipher( byte[] keyBytes ) throws  NoSuchAlgorithmException, 
																				NoSuchPaddingException, 
																				InvalidKeyException, 
																				InvalidAlgorithmParameterException, 
																				NoSuchProviderException {
		return generateDecryptingAesCipher( keyBytes, CTR_ZERO_VAL );
	}
	
	/**
	 * Returns an decrypting AES-CTR-Cipher with length of keyBytes and initial counter value ctr.
	 * 
	 * @param keyBytes The keyBytes for the secret key.
	 * @param ctr Counter value as 8 bytes data. ( big endian )
	 * @return A decrypting AES-CTR-Cipher with length of keyBytes and initial counter value ctr.
	 * 
	 * @throws NoSuchAlgorithmException
	 * @throws NoSuchPaddingException
	 * @throws InvalidKeyException
	 * @throws InvalidAlgorithmParameterException
	 * @throws NoSuchProviderException
	 */
	public static Cipher generateDecryptingAesCipher( byte[] keyBytes, byte[] ctr ) throws  NoSuchAlgorithmException, 
																							NoSuchPaddingException, 
																							InvalidKeyException, 
																							InvalidAlgorithmParameterException, 
																							NoSuchProviderException {
		
		// complete ctr key
		ctr = ArrayFunctions.merge( ctr, CTR_ZERO_VAL );
		
		IvParameterSpec spec = new IvParameterSpec( ctr );
		SecretKeySpec key = new SecretKeySpec( keyBytes, "AES" );
	    Cipher cipher = Cipher.getInstance( "AES/CTR/NoPadding" );

	    // encryption pass
	    cipher.init( Cipher.DECRYPT_MODE, key, spec );
		
		return cipher;
	}
	
	/**
	 * Appends the given data array to the message buffer. To complete the otr message
	 * the otr header ( ?OTR: ) and  otr footer ( . ) will be added.
	 * 
	 * @param message The buffer to append the otr message.
	 * @param data The otr message.
	 */
	public static void generateOTRMessage( StringBuffer message, byte[] data ) {
		message.append( OTR_MSG_HEADER ).append( new String( StringFunctions.encodeBase64( data ) )  ).append( "." );
	}

	/**
	 * Injects the OTR Error header to the given message.
	 * 
	 * @param message The buffer to inject the otr error headerr.
	 */
	public static void injectOTRErrorMessage( StringBuffer message ) {
		String info = message.toString();
		message.append( OTR_ERROR_HEADER ).append( info  );
	}

	
	/**
	 * Removes all whitespace-tags from this message buffer.
	 * 
	 * @param message The StringBuffer to clean.
	 */
	public static void removeWhitespaceTags( StringBuffer message ) {
		String tempString = message.toString();
		
		tempString = tempString.replaceAll( String.valueOf( WHITESPACE_TAG_BASIC ), "" );
		tempString = tempString.replaceAll( String.valueOf( WHITESPACE_TAG_V1 ), "" );
		tempString = tempString.replaceAll( String.valueOf( WHITESPACE_TAG_V2 ), "" );
		
		message.setLength( 0 );
		message.append( tempString );
	}
	
	/**
	 * Tests for valid protocol-version.
	 * 
	 * @param version The version to validate.
	 * @return <b>true</b> if version is valid, <b>false</b> otherwise.
	 */
	public static boolean isSupportedVersion( int version ) {
		switch( version ) {
			case PROTOCOL_VERSION_1:
			case PROTOCOL_VERSION_2:
				return true;
			default:
				return false;
		}
	}
	
	/**
	 * Returns the content between the otr header and the otr footer.
	 * 
	 * @param buffer The complete Message with header and footer.
	 * @return the The content between the otr header ( \?OTR: ) and the otr footer ( . ).
	 */
	public static String stripOtrMessage( StringBuffer buffer ) {
		
		String message = buffer.toString();
		
		int start = message.indexOf( OTR_MSG_HEADER );
		int end = message.indexOf( "." );
		
		if( start == -1 || end == -1)
			return null;
		
		start += OTR_MSG_HEADER.length();
		
		return buffer.substring( start, end );
	}
	
	/**
	 * Returns the content after the otr error header.
	 * 
	 * @param buffer The complete Message.
	 * @return the The content after the otr error header ( ?OTR Error: ).
	 */
	public static String stripOtrErrorMessage( StringBuffer buffer ) {
		
		String message = buffer.toString();
		
		int start = message.indexOf( OTR_ERROR_HEADER );
		
		if( start == -1 )
			return null;
		
		start += OTR_ERROR_HEADER.length();
		
		return buffer.substring( start, buffer.length() );
	}

	/**
	 * Checks the message for the Query-header and returns best available version.
	 * 
	 * @return Best available version. 0 is no version is available.
	 */
	public static int checkForOTRQueryHeader( StringBuffer message ) {
		
		boolean v1 = false;
		boolean v2 = false;
		
		String msg = message.toString();
		
		if(  msg.contains( OTR_QUERY_HEADER1 ) ) {
			v1 = true;

			int start = msg.indexOf( OTR_QUERY_HEADER1 ) + OTR_QUERY_HEADER1.length(); 
			int end = msg.indexOf( "?", start );
			
			String versions = msg.substring( start, end );
			
			if( versions.contains( "2" ) );	
				v2 = true;
			
		} else if ( msg.contains( OTR_QUERY_HEADER2 ) ) {

			int start = msg.indexOf( OTR_QUERY_HEADER2 ) + OTR_QUERY_HEADER2.length(); 
			int end = msg.indexOf( "?", start );
				
			String versions = msg.substring( start, end );
			
			if( versions.contains( "2" ) );	
				v2 = true;

		} else {
			return 0;
		}
		
		if( v2 )
			return JotrFunctions.PROTOCOL_VERSION_2;
		if( v1 )
			return JotrFunctions.PROTOCOL_VERSION_1;
		
		return 0;
	}
	
	/**
	 * Checks the message for the OTR error-header.
	 * 
	 * @return <b>true</b> if message contains "?OTR Error:", <b>false</b> otherwise.
	 */
	public static boolean checkForOTRErrorHeader( StringBuffer message ) {
		return message.toString().contains( OTR_ERROR_HEADER );
	}
	
	/**
	 * Checks the message for the OTR header.
	 * 
	 * @return <b>true</b> if message contains  "?OTR:", <b>false</b> otherwise.
	 */
	public static boolean checkForOTRHeader( StringBuffer message ) {
		return message.toString().contains( OTR_MSG_HEADER );
	}
	
	/**
	 * Checks the given String for known whitespace tags.
	 * 
	 * @return Returns the version of the found whitespacetag or -1 on error.
	 * @see de.karnik.libjotr.JotrConstants#PROTOCOL_VERSION_1
	 * @see de.karnik.libjotr.JotrConstants#PROTOCOL_VERSION_2
	 */
	public static int checkForWhitespaceTag( StringBuffer message ) {
		if( message.toString().contains( String.valueOf( WHITESPACE_TAG_BASIC ) ) ) {
			if ( message.toString().contains( String.valueOf( WHITESPACE_TAG_V2 ) ) ) {
				LogFunctions.log( "JotrFunctions.checkForWhitespaceTag", "JOTR:  contains v2 tag ");
				return PROTOCOL_VERSION_2;
			} else if( message.toString().contains( String.valueOf( WHITESPACE_TAG_V1 ) ) ) {
				LogFunctions.log( "JotrFunctions.checkForWhitespaceTag", "JOTR:  contains v1 tag ");
				return PROTOCOL_VERSION_1;
			}
		}
		return -1;
	}
	
	/**
	 * Injects the whitespace tag to the message. 
	 * 
	 * @param message The message to inject the tags.
	 */
	public static void injectWhitspaceTag( StringBuffer message ) {
			message.append( WHITESPACE_TAG_BASIC ).append( WHITESPACE_TAG_V1 )
				.append( WHITESPACE_TAG_V2 );
	}	
	
	/**
	 * Signs the given data with the DSA-key of the given JotrContext object.
	 * 
	 * @param context The context with the DSA-key to sign.
	 * @param data The data to sign.
	 * @return The DSA signature of the given data.
	 */
	public static byte[] signWithDSASig( JotrContext context, byte[] data ) {
		
		JotrDsaKeyStorage ks = context.getDsaKeyStorage();
		
        // construct the BC objects from pub key specs
        DSAParameters bcDSAParams = new DSAParameters( 	ks.getP().getValueAsBigInteger(),
        												ks.getQ().getValueAsBigInteger(),
        												ks.getG().getValueAsBigInteger() );
        
        DSAPrivateKeyParameters dsaPrivParms = new DSAPrivateKeyParameters( ks.getX().getValueAsBigInteger(), bcDSAParams );
        
        // and now do the signature verification
        DSASigner dsaSigner = new DSASigner();
        dsaSigner.init( true, dsaPrivParms );
        
        BigInteger[] rs = dsaSigner.generateSignature( data );
        
        if( rs.length == 2 ) {
        	byte[] r = ArrayFunctions.serialize( rs[ 0 ].toByteArray() );
        	byte[] s = ArrayFunctions.serialize( rs[ 1 ].toByteArray() );
        	
        	return ArrayFunctions.merge( r, s );
        }
		
		return null;
	}
	
	/**
	 * Verifies the given signature against the DSA-key and signed data.
	 * 
	 * @param context The context with the DSA-key to  verify against.
	 * @param data The data to verify against.
	 * @param r The r part of the signature.
	 * @param s The r part of the signature.
	 * @return <b>true</b> if verification was successful, <b>false</b> otherwise.
	 * @throws NoSuchAlgorithmException
	 */
	public static boolean verifyDSASig( JotrContext context, byte[] data, byte[] r, byte[] s ) throws NoSuchAlgorithmException {
		
		JotrDsaKeyStorage ks = context.getDsaKeyStorage();
		
        // construct the BC objects from pub key specs
        DSAParameters bcDSAParams = new DSAParameters( 	ks.getRemoteP().getValueAsBigInteger(),
        												ks.getRemoteQ().getValueAsBigInteger(),
        												ks.getRemoteG().getValueAsBigInteger() );
        
        DSAPublicKeyParameters dsaPubParms = new DSAPublicKeyParameters( ks.getRemoteY().getValueAsBigInteger(), bcDSAParams );
        
        // and now do the signature verification
        DSASigner dsaSigner = new DSASigner();
        dsaSigner.init( false, dsaPubParms );
        
        return dsaSigner.verifySignature( data, new BigInteger( 1, r ), new BigInteger( 1, s ) );
	}
	
	/**
	 * Computes the SSID from the given shared secret. 
	 * 
	 * @param sharedSecred The shared secret to calculate the SSID from.
	 * @return The SSID.
	 */
	public static byte[] computeSsid( JotrMpi sharedSecred ) {
		return computeAKEOtrKeys( SSID, sharedSecred );
	}
	
	/**
	 * Computes the specified AKE-OTR-key from the given shared secret.
	 * 
	 * @param selection The specification for the key.
	 * @param sharedSecret The shared secret to calculate the AKE-OTR-key from.
	 * @return The AKE-OTR-key.
	 * 
	 * @see de.karnik.libjotr.JotrConstants#BYTE_C
	 * @see de.karnik.libjotr.JotrConstants#BYTE_M1
	 * @see de.karnik.libjotr.JotrConstants#BYTE_M2
	 * @see de.karnik.libjotr.JotrConstants#BYTE_M1_I
	 * @see de.karnik.libjotr.JotrConstants#BYTE_M2_I
	 */
	public static byte[] computeAKEOtrKeys( int selection, JotrMpi sharedSecret ) {

		byte[] sharedSec = sharedSecret.getValueAsByteArray();
		
		switch( selection ) {
			case SSID:
				return ArrayFunctions.seperateDataFromArray( 0, 8, h2( (byte)SSID, sharedSec ) );
			case KEY_C:
				return ArrayFunctions.seperateDataFromArray( 0, 16, h2( BYTE_C, sharedSec ) );
			case KEY_C_I:
				return ArrayFunctions.seperateDataFromArray( 16, 16, h2( BYTE_C, sharedSec ) );
			case KEY_M1:
				return h2( BYTE_M1, sharedSec );
			case KEY_M2:
				return h2( BYTE_M2, sharedSec );
			case KEY_M1_I:
				return h2( BYTE_M1_I, sharedSec );
			case KEY_M2_I:
				return h2( BYTE_M2_I, sharedSec );
		}
		
		return null;	
	}
	
	/**
	 * Computes the sending or receiving aes key for the given shared secret.
	 * 
	 * @param b The parameter to specifiy the type.
	 * @param sharedSecret The shared secret to calculate the key from.
	 * @return The key specified by b.
	 * 
	 * @see de.karnik.libjotr.JotrConstants#KEY_C
	 * @see de.karnik.libjotr.JotrConstants#KEY_C_I
	 */
	public static byte[] computeSendRecvAesKey( byte b, JotrMpi sharedSecret ) {
		byte[] sharedSec = sharedSecret.getValueAsByteArray();
		byte[] retVal = h1( b, sharedSec );
		
		return ArrayFunctions.seperateDataFromArray( 0 , 16, retVal ); 
	}

	/**
	 * Calculates and returns the sha1 hash of the concatenation of b and input.
	 * 
	 * sha1( b + input )
	 * 
	 * @param b The byte to hash.
	 * @param input The data to hash.
	 * @return The sha1 hash of the concatenation of b and input.
	 */
	private static byte[] h1( byte b, byte[] input ) {
		byte[] retVal = new byte[ input.length + 1 ];
		retVal[ 0 ] = b;
		
		for( int i = 1; i < retVal.length; i++ )
			retVal[ i ] = input[ i - 1 ];

		return HashFunctions.sha1( retVal );
	}
	
	/**
	 * Calculates and returns the sha256 hash of the concatenation of b and input.
	 * 
	 * sha256( b + input )
	 * 
	 * @param b The byte to hash.
	 * @param input The data to hash.
	 * @return The sha256 hash of the concatenation of b and input.
	 */
	private static byte[] h2( byte b, byte[] input ) {
		byte[] retVal = new byte[ input.length + 1 ];
		retVal[ 0 ] = b;
		
		for( int i = 1; i < retVal.length; i++ )
			retVal[ i ] = input[ i - 1 ];

		return HashFunctions.sha256( retVal );
	}
	
}
